from dataclasses import dataclass
from typing import Tuple, Any
from type import EnumViewType
import numpy as np


@dataclass
class PhysicalProperties:
    first_position: Tuple[float, float, float]  # 患者坐标系下第一张影像的位置。
    last_position: Tuple[float, float, float]  # 患者坐标系下最后张影像的位置。
    image_orientation: Tuple[float, float, float, float, float, float]  # 方向余弦
    direction_matrix: Any

    def set_physical_property(self, slices):
        self.first_position = slices[0].ImagePositionPatient
        self.last_position = slices[-1].ImagePositionPatient
        self.image_orientation = slices[0].ImageOrientationPatient

        orientation = np.array(slices[0].ImageOrientationPatient).reshape(2, 3)
        row_vector = orientation[0]
        col_vector = orientation[1]
        slice_vector = np.array(self.last_position) - np.array(
            self.first_position)
        slice_vector = slice_vector / np.linalg.norm(slice_vector)
        direction_matrix = np.vstack([row_vector, col_vector, slice_vector]).T
        self.direction_matrix = direction_matrix


class CoordTransfer:
    # 物理坐标 => Origin_voxel => Norm_voxel => screen
    # screen => norm_voxel => origin_voxel => 物理坐标
    def __init__(self):
        self.spacing = (1, 1, 1)
        self.physical_properties = PhysicalProperties

    def init(self, slices):
        self.spacing = np.array([
            float(slices[0].PixelSpacing[0]),
            float(slices[0].PixelSpacing[1]),
            float(slices[0].get('SliceThickness')),
        ])
        self.physical_properties.first_position = slices[0].ImagePositionPatient
        self.physical_properties.last_position = slices[-1].ImagePositionPatient
        self.physical_properties.image_orientation = slices[0].ImageOrientationPatient

        orientation = np.array(slices[0].ImageOrientationPatient).reshape(2, 3)
        row_vector = orientation[0]
        col_vector = orientation[1]
        slice_vector = np.array(self.physical_properties.last_position) - np.array(
            self.physical_properties.first_position)
        slice_vector = slice_vector / np.linalg.norm(slice_vector)
        direction_matrix = np.vstack([row_vector, col_vector, slice_vector]).T
        self.physical_properties.direction_matrix = direction_matrix

    # 物理坐标 => Origin_voxel
    def physical_to_origin_voxel(self, physical_point):
        relative_physical_point = np.array(physical_point) - self.physical_properties.first_position
        inverse_direction_matrix = np.linalg.inv(self.physical_properties.direction_matrix)
        voxel_point = inverse_direction_matrix @ relative_physical_point
        voxel_point_rounded = np.round(voxel_point / self.spacing)
        return tuple(voxel_point_rounded.astype(int))

    # origin_voxel => screen
    def origin_voxel_to_screen(self, view_type, origin_voxel_point, transform_matrix):
        norm_voxel_point = self.origin_voxel_to_norm_voxel_vector2(view_type, origin_voxel_point)
        x, y = norm_voxel_point
        point_homogeneous = np.array([y, x, 1])
        transformed_point = transform_matrix @ point_homogeneous
        return round(transformed_point[1]), round(transformed_point[0])

    # origin => physical
    def origin_voxel_to_physical(self, voxel_point):
        physical_point = (
                self.physical_properties.first_position +
                self.physical_properties.direction_matrix @ (voxel_point * self.spacing)
        )
        physical_point_rounded = tuple(float(round(each, 3)) for each in physical_point)
        return physical_point_rounded

    def screen_to_origin_voxel(self, view_type, transform_matrix, vector2_point):
        point_homogeneous = np.array([vector2_point[1], vector2_point[0], 1])
        inverse_matrix = np.linalg.inv(transform_matrix)
        norm_point = inverse_matrix @ point_homogeneous
        origin_point = self.norm_voxel_to_origin_voxel_vector2(view_type, norm_point)
        return origin_point

    def origin_voxel_to_norm_voxel_vector2(self, view_type, origin_voxel):
        spacing = self.spacing
        norm_point = None
        match view_type:
            case EnumViewType.Stack.value:
                norm_point = origin_voxel
            case EnumViewType.Sag.value:
                norm_point = (origin_voxel[0] * spacing[0], origin_voxel[1] * spacing[2])
            case EnumViewType.Cor.value:
                norm_point = (origin_voxel[0] * spacing[0], origin_voxel[1] * spacing[2])
            case EnumViewType.Ax.value:
                norm_point = (origin_voxel[0] * spacing[0], origin_voxel[1] * spacing[1])
        return norm_point

    def norm_voxel_to_origin_voxel_vector2(self, view_type, norm_voxel):
        spacing = self.spacing
        origin_voxel = None
        match view_type:
            case EnumViewType.Stack.value:
                origin_voxel = norm_voxel
            case EnumViewType.Sag.value:
                origin_voxel = (norm_voxel[0] / spacing[2], norm_voxel[1] / spacing[0])
            case EnumViewType.Cor.value:
                origin_voxel = (norm_voxel[0] / spacing[2], norm_voxel[1] / spacing[0])
            case EnumViewType.Ax.value:
                origin_voxel = (norm_voxel[0] / spacing[0], norm_voxel[1] / spacing[1])
        return origin_voxel

    def origin_voxel_to_norm_voxel_vector3(self, origin_voxel):
        norm_voxel = (round(origin_voxel[0] * self.spacing[2]),
                      round(origin_voxel[1] * self.spacing[0]),
                      round(origin_voxel[2] * self.spacing[1]))
        return norm_voxel
